from sympy import symbols, nroots, simplify, re, im, Poly
from numpy.polynomial import Polynomial as npPoly
import numpy as np
import re

def solveEq(eq) -> dict:
  λ = symbols("λ")
  # eq = 117-56*λ+7*λ*λ-λ*λ*λ
  # eq = "-λ*λ*λ+15*λ*λ+18*λ"
  #eq = "-λ*(1-λ)*(1-λ)"
  eq = "-35666*λ**3+388*λ**4-λ**5+0+0"
  # eq = simplify(eq)
  # eq = simplify("227.14965502300004+4.557*λ*λ+35.767477*λ-λ*λ*λ")

  solutions = nroots(eq)
  threshold = 1e-9
  coe = Poly(eq, λ).all_coeffs()
  print(coe)
  input("wait")
  solutions = [re(sol) if abs(im(sol)) < threshold else sol for sol in solutions]
  # solutions = [simplify(sol) for sol in solutions]
  if all(s.is_real for s in solutions):
    allRealRoots = True  
    solutions = [float(str(s.evalf(10))) for s in solutions]
  else:
    allRealRoots = False
    solutions = [str(s.evalf(3)) for s in solutions]
  print({"roots":solutions,"allRealRoots":allRealRoots})
  return {"roots":solutions,"allRealRoots":allRealRoots}

def eqSimplify(eq)-> dict:
  eq = "(-21*λ-13*λ*λ+17*λ*λ*λ+18-λ*λ*λ*λ)*λ/((1-λ)*(1-λ))"
  eq = "(-21*λ-13*λ*λ+17*λ*λ*λ+18-λ*λ*λ*λ)*λ/((1-λ)*(1-λ))"
  eq = simplify(eq)
  print(eq)

def solveWithNumpy(eq)-> dict:
  coef = []
  eq = "28.921957+10.518*λ+λ*λ"
  eq = "+" + eq if not eq.startswith("-") else eq
  l = re.split(r"([\+\-])",eq)
  l = ["1*"+l[i] if l[i].startswith("λ") else l[i] for i in range(len(l))]
  l = list(filter(lambda x: x != "",l))
  l = [l[i]+l[i+1] for i in range(0, len(l)-1,2)]
  m = [len(re.findall(r"\*λ",l)) for l in l]
  l= [float(l.replace(r"*λ","")) for l in l]
  for i in range(m[-1]+1):
    if i in m:
      coef.append(l[m.index(i)])
    else:
      coef.append(0)
  solutions = npPoly(coef,symbol="λ").roots()
  if (all(np.isreal(solutions))):
    allRealRoots = True  
    solutions = [round(s.astype(float),10) for s in solutions]
  else:
    allRealRoots = False
    solutions = np.array(solutions).astype(str)
  print({"roots":solutions,"allRealRoots":allRealRoots})
  return {"roots":solutions,"allRealRoots":allRealRoots}



if __name__ == "__main__":
  # solveEq(0)
  # eqSimplify(0)
  solveWithNumpy(0)
